#!/usr/bin/env python
##
#  Project: Kelindo - Unsupervised Music Organizer  
#  Author: Nicola Bicocchi <nicola.bicocchi@gmail.com>, Nicolo' Torreggiani <nicolo.torreggiani@gmail.com>, Gianluca Puglia <gianluca.puglia@gmail.com>, Matteo Senardi <pualien@gmail.com>
#  Copyright: 2013 Nicola Bicocchi, Nicolo' Torreggiani, Gianluca Puglia, Matteo Senardi
#  License: GPL-2+
#  This program is free software; you can redistribute it and/or modify it
#  under the terms of the GNU General Public License as published by the Free
#  Software Foundation; either version 2 of the License, or (at your option)
#  any later version.
# 
#  This program is distributed in the hope that it will be useful, but WITHOUT
#  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
#  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
#  more details.
# 
# On Debian GNU/Linux systems, the full text of the GNU General Public License
# can be found in the file /usr/share/common-licenses/GPL-2.
##

import re
import os
import os.path
import wave
import contextlib
import shutil
import subprocess

from mutagen.mp3 import MP3
from mutagen.id3 import TIT2, TALB, TPE1, TRCK, TDRC, TCON, COMM
from mutagen.flac import FLAC
from mutagen.mp4 import MP4
from mutagen.asf import ASF
from mutagen.musepack import Musepack
from mutagen.oggvorbis import OggVorbis

from kelindo.core.album import Album
from kelindo.core.musicbrainz import MusicBrainz
from kelindo.core.musicbrainz import MBError
from kelindo.core.bagofwords import BagOfWords
from kelindo.core.observable import Observable

class SimilarityError(Exception): pass

class Kelindo(Observable):

    CODES = ['OK', 'NO_MEDIA', 'TOOMANY_MEDIA', 'NO_TNUMBERS', 'NO_DURATIONS', 'UNSUPPORTED', 'NO_MATCH', 'MATCH']

    def __init__(self, path):
        super(Kelindo, self).__init__()
        
        self.dir = path
        self.__invalidExts = set()
        self.__validExts = ['.mp3', '.flac', '.wav', '.ogg', '.mp4', '.m4a', '.asf', '.mpc']
        
        self.albums = []
        
        self.mb = MusicBrainz()
        self.stat_start = None
        self.stat_stop = None
        
                
    def __coherent_track_position(self, album):
        """ Checks if track numbers within an album are coherent """
        album.tracks = sorted(album.tracks, key=lambda track: track['trackposition'])
        for i in range(len(album.tracks)):
            if album.tracks[i]['trackposition'] != i + 1:
                return False
        return True
        
        
    def __build_album(self, root, files):
        """ Roughly classifies dir-tree """     
        album = Album(root) 
        exts = set()

        for file in files:
            trackName, ext = os.path.splitext(file)

            if ext.lower() in self.__validExts:
                exts.add(ext)
                track = {
                    'filename' : file,
                    'trackposition' : 0,
                }
                album.tracks.append(track)
                
                # search track numbers
                match = re.search('(\d{1,2})', file) 
                if match:
                    track['trackposition'] = int(match.groups()[0])
            # debug only: annotate unrecognized extensions
            else:
                self.__invalidExts.add(ext)
                
        if len(exts) == 1:
            album.format = exts.pop().lower()
            album.format = album.format.replace('.','')
        else:
            album.format = 'misc'
                
        # empty folder
        if (len(album.tracks) < 1):
            album.status = 'NO_MEDIA'
            return album
        
        # too many files to be an audio cd
        if (len(album.tracks) > 99):
            album.status = 'TOOMANY_MEDIA'
            return album
        
        try:
            if not self.__coherent_track_position(album):
                self.__get_metadata(album)
            else:
                self.__get_metadata(album, ['TALB', 'TIT2', 'TPE1', 'TCON'])
            
            if not self.__coherent_track_position(album):
                album.status = 'NO_TNUMBERS'
                return album
        except:
            album.status = 'UNSUPPORTED'
            return album
            
        album.status = 'OK'
        return album
    
    
    def __get_metadata(self, album, tags = ['TALB', 'TIT2', 'TPE1', 'TCON', 'TRCK']):
        """ Extracts metadata for all tracks of an album """
        tagMap = {
                'TALB': 'albumtitle',
                'TIT2': 'tracktitle',
                'TPE1': 'trackauthors',
                'TCON': 'trackgenre',
                'TRCK': 'trackposition',
        }

        for track in album.tracks:
            filename = os.path.join(album.dir, track['filename'])
            extension = os.path.splitext(track['filename'])[1].lower()
            
            if extension == '.mp3':
                meta = MP3(filename)
            elif extension == '.flac':
                meta = FLAC(filename)
            elif extension == '.mp4' or extension.lower() == '.m4a':
                meta = MP4(filename)
            elif extension == '.asf':
                meta = ASF(filename)
            elif extension == '.mpc':
                meta = Musepack(filename)
            elif extension == '.ogg':
                meta = OggVorbis(filename)
            elif extension == '.wav':
                album.need_conversion = True
                with contextlib.closing(wave.open(filename,'r')) as wav:
                    track['tracklength'] = wav.getnframes() / wav.getframerate()
            else:
                raise IOError
                
            try:
                track['tracklength'] = int(meta.info.length) 
                track['trackbitrate'] = (float(os.path.getsize(filename) * 8.0) / 1000.0) / track['tracklength']
            except AttributeError:
                pass

            for tag in tags:
                try:
                    track[tagMap[tag]] = meta[tag].text[0]
                except KeyError:
                    pass
                    
            if 'TRCK' in tags:
                track['trackposition'] = int(unicode(track['trackposition']).split('/')[0])

        album.avg_bitrate = 0.0
        for track in album.tracks:
            album.avg_bitrate += track['trackbitrate']
        album.avg_bitrate /= len(album.tracks)
        return
                
                
    def __get_best_match(self, album, candidates):
        """ Returns the best match among a set of candidates """
        album_bag = BagOfWords()
        for word in album.dir.split('/')[-3:]: album_bag.add_string(word)
        for track in album.tracks:
            for field in ['albumtitle', 'tracktitle', 'trackauthors', 'filename']:
                try:
                    album_bag.add_string(track[field])
                except KeyError:
                    pass
            
        best_match = [0.0, None]
        for candidate in candidates:
            if len(album.tracks) != len(candidate):
                continue
            
            other_bag = BagOfWords()
            for track in candidate:
                for field in ['albumtitle', 'albumauthors', 'tracktitle', 'trackauthors']:
                    other_bag.add_string(track[field])

            similarity = album_bag.similarity(other_bag)
            if (similarity >= best_match[0]):
                best_match[0] = similarity
                best_match[1] = candidate
                
        if best_match[1] is None:
            raise SimilarityError
        return best_match
            
            
    def __classify_album(self, album): 
        #TODO: get_MBID and get_albums might fail
        # we still do not handle that case
        MBIDs = self.mb.get_MBID(album.get_durations())
        try:    
            candidates = []
            for MBID in MBIDs:
                candidate = self.mb.get_albums(MBID)
                candidates.extend(candidate)
            album.distance, album.best_candidate = self.__get_best_match(album, candidates)
            album.status = 'MATCH'
            album.finalize = True
        
        except MBError:
            album.status = 'NO_MATCH'
            return
        except SimilarityError:
            album.status = 'NO_MATCH'
            return
        return
        
        
    def __write_tags(self, dst, track, ext):
        if ext == '.mp3':
            audio = MP3(dst)
            audio['TIT2'] = TIT2(encoding=3, text=[track['tracktitle']])
            audio['TALB'] = TALB(encoding=3, text=[track['albumtitle']])
            audio['TRCK'] = TRCK(encoding=3, text=[track['trackposition']])
            audio['TPE1'] = TPE1(encoding=3, text=[track['albumauthors']])
            audio['TDRC'] = TDRC(encoding=3, text=[track['albumyear']])
            #audio['TCON'] = TCON(encoding=3, text=[track['trackgenre']])
            audio['COMM'] = COMM(encoding=3, text=['Beautyfied by Kelindo!'])
        else:
            if ext == '.flac':
                audio = FLAC(dst)
            elif ext == '.mp4' or ext.lower() == '.m4a':
                audio = MP4(dst)
            elif ext == '.asf':
                audio = ASF(dst)
            elif ext == '.mpc':
                audio = Musepack(dst)
            elif ext == '.ogg':
                audio = OggVorbis(dst)
            else:
                pass
            
            audio['title'] = track['tracktitle']
            audio['album'] = track['albumtitle']
            audio['tracknumber'] = track['trackposition']
            audio['artist'] = track['albumauthors']       
            audio['date'] = track['albumyear']
            #audio['genre'] = track['trackgenre']
            audio['comment'] = 'Beautyfied by Kelindo!'
        audio.save()
        return
        
        
    def __convert(self, src, dst, fmt, avg_bitrate):
        def closest(target, collection) :
            return min((abs(target - i), i) for i in collection)[1]
        
        bitrate = closest(avg_bitrate, [128, 256, 320])        
        if fmt == 'mp3':
            subprocess.call(['sox', '--multi-threaded', src, '--compression', 
                             '%d.2' % (bitrate,), dst])
        elif fmt == 'ogg':
            ogg_compression = {128: '5', 256 : '7', 320 : '9'}
            subprocess.call(['sox', '--multi-threaded', src, 
                             '--compression', ogg_compression[bitrate], dst])
        else:
            subprocess.call(['sox', '--multi-threaded', src, dst])
        return
    
    def __copy(self, src, dst):
        shutil.copyfile(src, dst)
        return
    
    def __move(self, src, dst):
        shutil.move(src, dst)
        return
        
    def __write_log(self, filename):
        outlist = []
        for album in self.albums:
            albumtitle = albumauthors = ''
            
            if album.status == 'MATCH':
                albumauthors = album.best_candidate[0]['albumauthors']
                albumtitle = album.best_candidate[0]['albumtitle']
                
            outlist.append('%s||%s||%s||%s' % (
                    album.dir, 
                    album.status, 
                    albumauthors, 
                    albumtitle))
                    
        f = open(filename, 'w')
        f.write('\n'.join(outlist))
        f.close()
        return
        
        
    def find_candidate_albums(self):
        self.fire(key='candidate_started')
        for root, dirs, files in os.walk(self.dir):
            album = self.__build_album(root, files) 
            self.albums.append(album)
            self.fire(key='candidate_found')
            
        self.stat_start = self.get_stats() 
        if self.stat_start['OK'] == 0:
            self.fire(key='candidate_no')
        else:
            self.fire(key='candidate_done')
        return
        
        
    def classify_albums(self):
        album_total = self.get_stats()['OK']
        album_done = 0
        self.fire(key='album_started', album=None, total=album_total, done=album_done)
        for album in self.albums:
            if album.status == 'OK':
                self.__classify_album(album)
                album_done += 1
                self.fire(key='album_found', album=album, total=album_total, done=album_done)
        
        self.stat_stop = self.get_stats()
        self.fire(key='album_done', album=None, total=album_total, done=album_done)
        return
        
        
    def finalize_albums(self, root_dst_dir, fmt, move_folder):
        root_dst_dir = os.path.join(root_dst_dir, "Kelindo")

        tracks_total = tracks_done = 0
        
        for album in self.albums:
            if album.finalize == True:                
                if fmt == 'as source':
                    if album.format == 'misc':
                        album.new_format = 'mp3'
                        album.need_conversion = True
                    elif album.format == 'wav':
                        album.new_format = 'flac'
                        album.need_conversion = True
                else:
                    if not album.format == fmt:
                        album.new_format = fmt
                        album.need_conversion = True    
                tracks_total += album.get_track_number()
        
        self.fire(key='finalize_started', album=None, total=tracks_total, done=tracks_done)
        
        for album in self.albums:
            if album.status != 'MATCH' or album.finalize != True:
                continue
            
            dst_dir = os.path.join(root_dst_dir, album.best_candidate[0]['albumauthors'])
            if int(album.best_candidate[0]['trackmedium']) != 1:
                medium = ' (CD %s)' % album.best_candidate[0]['trackmedium']
                dst_dir = os.path.join(dst_dir, album.best_candidate[0]['albumtitle'] + medium)
            else:
                dst_dir = os.path.join(dst_dir, album.best_candidate[0]['albumtitle'])
            
            if not os.path.exists(dst_dir):
                os.makedirs(dst_dir)
                album.new_dir = dst_dir
            else:
                print 'Album already exists, overwrite', dst_dir
                shutil.rmtree(dst_dir)
                os.makedirs(dst_dir)
                album.new_dir = dst_dir

            for idx in range(1, album.get_track_number() + 1):                
                src_filename, src_ext = os.path.splitext(album.get_track(idx)['filename'])
                src = os.path.join(album.dir, src_filename + src_ext)  
                                
                dst_track = album.get_best_candidate_track(idx)
                dst_filename = '%s - %s' % (
                    dst_track['trackposition'].zfill(2),
                    dst_track['tracktitle'])
                dst = os.path.join(dst_dir, dst_filename)
                
                if album.need_conversion:
                    dst = dst + '.' + fmt
                    self.__convert(src, dst, fmt, album.avg_bitrate)
                    self.__write_tags(dst, dst_track, '.' + fmt)
                else:
                    dst = dst + src_ext
                    if move_folder: 
                        self.__move(src, dst)
                        os.chmod(dst, 0644)
                    else:
                        self.__copy(src, dst)
                        os.chmod(dst, 0644)
                    self.__write_tags(dst, dst_track, src_ext.lower())
                
                tracks_done += 1
                self.fire(key='finalize_track_done', album=album, total=tracks_total, done=tracks_done)
                
            if move_folder:
                shutil.rmtree(album.dir)
          
        self.__write_log(os.path.join(root_dst_dir, 'kelindo.log'))
        self.fire(key='finalize_done', album=None, total=tracks_total, done=tracks_done)
        return
        
        
    def get_stats(self):
        """ Returns some stats """
        stat = {}
        for status in Kelindo.CODES:
            stat[status] = 0
        
        for album in self.albums:
            stat[album.status] += 1
        return stat
        

        
        
